import paramiko
import threading
import time
import socket
from typing import List, Tuple

class SSHForwarder:
    def __init__(self):
        self.client = None
        self.forwardings = []
        self.is_connected = False
        self.tunnel_threads = []
        self.local_sockets = []  # 存储本地监听socket
        
    def connect(self, hostname: str, port: int, username: str, password: str) -> bool:
        """建立SSH连接"""
        try:
            self.client = paramiko.SSHClient()
            self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            self.client.connect(hostname=hostname, port=port, username=username, password=password)
            self.is_connected = True
            return True
        except Exception as e:
            print(f"SSH连接失败: {e}")
            return False
    
    def disconnect(self):
        """断开SSH连接"""
        if self.client:
            self.client.close()
            self.is_connected = False
            
    def add_port_forwarding(self, local_port: int, remote_host: str, remote_port: int):
        """添加端口转发规则"""
        forwarding = {
            'local_port': local_port,
            'remote_host': remote_host,
            'remote_port': remote_port
        }
        self.forwardings.append(forwarding)
        
    def start_forwarding(self) -> bool:
        """启动所有端口转发"""
        if not self.is_connected:
            return False
            
        try:
            # 为每个转发规则创建一个线程
            for forwarding in self.forwardings:
                local_port = forwarding['local_port']
                remote_host = forwarding['remote_host']
                remote_port = forwarding['remote_port']
                
                # 创建SSH隧道的线程
                thread = threading.Thread(
                    target=self._create_tunnel,
                    args=(local_port, remote_host, remote_port)
                )
                thread.daemon = True
                thread.start()
                self.tunnel_threads.append(thread)
            
            return True
        except Exception as e:
            print(f"端口转发启动失败: {e}")
            return False
            
    def _create_tunnel(self, local_port: int, remote_host: str, remote_port: int):
        """创建单个SSH隧道"""
        try:
            # 创建本地监听socket
            local_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            local_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            local_socket.bind(("127.0.0.1", local_port))
            local_socket.listen(5)
            
            print(f"成功创建本地监听端口: {local_port}")
            
            # 将socket添加到列表中，用于后续关闭
            self.local_sockets.append(local_socket)
            
            while True:
                try:
                    # 等待客户端连接
                    client_socket, addr = local_socket.accept()
                    print(f"收到新连接: {addr} -> 本地端口 {local_port}")
                    
                    # 创建处理线程
                    thread = threading.Thread(
                        target=self._handle_client,
                        args=(client_socket, remote_host, remote_port)
                    )
                    thread.daemon = True
                    thread.start()
                except Exception as e:
                    print(f"接受连接时出错: {e}")
                    break
                    
        except Exception as e:
            print(f"创建隧道失败: {e}")
    
    def _handle_client(self, client_socket, remote_host: str, remote_port: int):
        """处理客户端连接"""
        try:
            # 创建SSH通道连接远程主机
            transport = self.client.get_transport()
            channel = transport.open_channel("direct-tcpip", (remote_host, remote_port), ("127.0.0.1", 0))
            
            print(f"建立SSH隧道: 客户端 -> {remote_host}:{remote_port}")
            
            # 双向数据传输
            def forward_data(src, dst):
                try:
                    while True:
                        data = src.recv(4096)
                        if not data:
                            break
                        dst.send(data)
                except Exception as e:
                    print(f"数据转发错误: {e}")
                finally:
                    try:
                        src.close()
                        dst.close()
                    except:
                        pass
            
            # 启动双向传输线程
            thread1 = threading.Thread(target=forward_data, args=(client_socket, channel))
            thread2 = threading.Thread(target=forward_data, args=(channel, client_socket))
            
            thread1.daemon = True
            thread2.daemon = True
            
            thread1.start()
            thread2.start()
            
        except Exception as e:
            print(f"处理客户端连接失败: {e}")
            try:
                client_socket.close()
            except:
                pass
    
    def stop_forwarding(self):
        """停止所有端口转发"""
        # 关闭所有本地socket
        for sock in self.local_sockets:
            try:
                sock.close()
            except Exception as e:
                print(f"关闭socket时出错: {e}")
        
        # 停止所有线程
        for thread in self.tunnel_threads:
            if thread.is_alive():
                thread.join(timeout=1)
        
        self.forwardings.clear()
        self.tunnel_threads.clear()
        self.local_sockets.clear()